import numpy as np
import matplotlib.pyplot as plt
import torch
import data
import model


if __name__ == '__main__':
    # cifar10 classes
    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    train_loader, test_loader = data.get_cifar_data(32)
    data_iter = iter(test_loader)
    images, labels = next(data_iter)

    idx = 2
    image = images[idx].numpy()
    image = np.transpose(image, (1, 2, 0))
    plt.imshow(image)
    plt.show()
    print(classes[labels[idx].numpy()])

    image_batch = image.reshape(-1, 3, 32, 32)
    image_tensor = torch.from_numpy(image_batch)

    model = model.ConvNet(10)
    model_info = torch.load("model_info.ckpt")
    model.load_state_dict(model_info["model"])
    model.eval()
    output = model(image_tensor)
    _, predicted = torch.max(output.data, 1)
    pre = predicted.numpy()
    print(pre)
    print(classes[pre[0]])
